using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace JufGame
{
    public class LossFunc : BaseLayer
    {
		public ComputeBuffer lossBuffer;
		private ComputeBuffer targetBuffer;
		private int kernelComputeLoss;

        public LossFunc(ComputeShader shader, int batch, int inputCount) : base(shader, batch, inputCount)
        {
			kernelComputeLoss = shader.FindKernel("ComputeLoss");
			lossBuffer = new ComputeBuffer(inputCount, sizeof(float));
			outputBuffer = new ComputeBuffer(batch * inputCount, sizeof(float));
        }

        public void SetTarget(float[] target)
		{
			targetBuffer ??= new ComputeBuffer(target.Length, sizeof(float));
			targetBuffer.SetData(target);
		}

        public override void Forward(ComputeBuffer inputBuffer)
        {
			compute.SetInt("inputCount", inputCount);
			compute.SetInt("batch", batch);
			compute.SetBuffer(kernelForward, "inputBuffer", inputBuffer);
			compute.SetBuffer(kernelForward, "outputBuffer", outputBuffer);
			compute.Dispatch(kernelForward, threadGroupsX, threadGroupsY, 1);

			compute.SetBuffer(kernelComputeLoss, "outputBuffer", outputBuffer);
			compute.SetBuffer(kernelComputeLoss, "targetBuffer", targetBuffer);
			compute.SetBuffer(kernelComputeLoss, "lossBuffer", lossBuffer);
			compute.Dispatch(kernelComputeLoss, threadGroupsX, threadGroupsY, 1);
        }

        public override void Backward(ComputeBuffer dOutputBuffer)
        {
			compute.SetInt("inputCount", inputCount);
			compute.SetInt("batch", batch);
			dInputBuffer ??= new ComputeBuffer(batch * inputCount, sizeof(float));
			compute.SetBuffer(kernelBackward, "outputBuffer", outputBuffer);
			compute.SetBuffer(kernelBackward, "targetBuffer", targetBuffer);
			compute.SetBuffer(kernelBackward, "dInputBuffer", dInputBuffer);
			compute.Dispatch(kernelBackward, threadGroupsX, threadGroupsY, 1);
        }

        public override void Release()
        {
            base.Release();
			targetBuffer?.Release();
			lossBuffer?.Release();
        }
    }
}
